!==========================================================================
 module data
   implicit none
   real (kind=8), dimension(:), allocatable :: xgrid, zgrid, delx, delz
   real (kind=8), dimension(:,:), allocatable :: u, v
!==========================================================================

 end module data

!==================================================================================================
 program stream
! The program reads grid and velocity from a YACC output file and either integrates the velocity
! via 2nd order Runge-Kutta to produce streamlines (default) or computes the
! streamfunction (default=no).
!
! The program needs 5 input parameters and is started as follows
!
! ./stream '/outputdir/' timestep dt nsource nsteps
!
! where
! - /outputdir/ is the path of the directory containing the output, from which
!   the grid and the velocity field are read
! - timestep is the time-step written in the ususal format, e.g. 000100
! - dt is the time-interval of particle advection (typically ~1d-5)
! - nsource is the number of source points (generated randomly) used to
!   start a streamline the more points the more streamlines you'll see 
!   (typically from 100 to a few hundred)
! - nsteps is the number of time steps to perform for each starting point (the
!   larger, the longer is the streamline plotted. Typically from 100 to 1000)
!==================================================================================================
   use data
   implicit none
   integer :: i,j,k,p,q,nsource,nsteps,nx,nz
   logical, parameter :: streamlines=.true., streamfunction=.false.
   real (kind=8) :: a,s1,s2,psi,x0,z0,x1,z1,dt
   character (len=100) :: ch, dirout, tstep, arg
!==============================================================

!------------------------------
! Read input data from screen
!------------------------------
   do i = 1, iargc()
     call getarg(i, arg)
     if(i.eq.1) dirout=arg
     if(i.eq.2) tstep=arg
     if(i.eq.3) read(arg,*) dt
     if(i.eq.4) read(arg,*) nsource
     if(i.eq.5) read(arg,*) nsteps
   end do
   
!------------------------
! Read yacc output file
!------------------------
   open(1,file=trim(dirout)//'grd_'//trim(tstep)//'_.dat')
   read(1,*) ch,ch,i,ch,a,ch,nx,ch,nz
   nx=nx-2
   nz=nz-2
   allocate(xgrid(0:nx+1), zgrid(0:nz+1), delx(1:nx+1), delz(1:nz+1), u(0:nx+1,0:nz+1), v(0:nx+1,0:nz+1))

   do i=0,nx+1
     do j=0,nz+1
       read(1,*) xgrid(i),zgrid(j),u(i,j),v(i,j)
     end do
   end do
   do i=1,nx+1
     delx(i)=xgrid(i)-xgrid(i-1)
   end do
   do j=1,nz+1
     delz(j)=zgrid(j)-zgrid(j-1)
   end do
   close(1)

!--------------------------------------------------------------   
! Compute streamlines via advection of the velocity field
!--------------------------------------------------------------   
   if(streamlines) then

   open(1,file='streamlines_'//trim(tstep)//'_.dat')
   do i=1,nsource
       a=dble(i)
       call random_number(a)
       x0=a

       a=dble(i+10)
       call random_number(a)
       z0=a

       do k=1,nsteps
         call rk2(x0,z0,dt,x1,z1)
         write(1,99) x1, z1
         x0=x1
         z0=z1
       end do
! Write a tag that is recognized by GMT to separate streamlines        
       write(1,*)'>'

   end do
   close(1)

   endif

!--------------------------------------------------------------   
! Compute the streamfunction (by default no)
!--------------------------------------------------------------   
   if(streamfunction) then

   open(1,file='streamfunc_'//trim(tstep)//'_.dat')
   psi=0d0
   do i=1,nx+1
   
     s1=0d0
     do p=1,i
       s1=s1+v(p,0)*(xgrid(p)-xgrid(p-1))
     end do

     do j=1, nz+1

       s2=0d0
       do q=1,j-1
         s2=s2+u(i,q)*(zgrid(q)-zgrid(q-1))
       end do
      
       psi=s2-s1 
       write(1,99) xgrid(i)-(xgrid(i)-xgrid(i-1))/2, zgrid(j)-(zgrid(j)-zgrid(j-1))/2, psi

     end do
   end do
   close(1)

   endif

99 format(2E16.4)


 contains

!===================================================================
 subroutine rk2(xold,zold,dt,xnew,znew)
! 2nd-order in space/1st order in time Runge-Kutta to
! find the new position (xnew,znew) of a marker initially 
! at (xold,zold) passively advected by the flow (u,v)
!===================================================================
   implicit none
   real (kind=8), intent(in) :: xold,zold,dt
   real (kind=8), intent(out) :: xnew,znew
   real (kind=8) :: x1,x2,z1,z2,k1,k2
!===================================================================
   
   x1=xold; z1=zold
   call grid2mk(x1,z1,u,k1)
   x2=xold+dt/2*k1
   call grid2mk(x2,z1,u,k2)
   xnew=xold+dt*k2

   x1=xold; z1=zold
   call grid2mk(x1,z1,v,k1)
   z2=zold+dt/2*k1
   call grid2mk(x1,z2,v,k2)
   znew=zold+dt*k2

 end subroutine rk2   

!=====================================================================================
 subroutine grid2mk(xm,zm,grdval,mkval)
! Given the grid, find the four nodes that surround (xm,zm) 
! and compute the value 'mkval' via linear interpolation from grid values 'grdval' 
!
! Input: 
! xm => real. x-coord of the interpolation point
! zm => real. z-coord of the interpolation point
! grdval => real, array(0:nx+1,0:nz+1). Field to be interpolated
! Output:
! mkval => real. Interpolated value at (xm,zm)
!=====================================================================================
   use data
   implicit none
   integer :: i1,i2,j1,j2
   real (kind=8), intent(in) :: xm,zm
   real (kind=8), intent(out) :: mkval
   real (kind=8) :: dxm,dzm
   real (kind=8), dimension(0:,0:) :: grdval
!===========================================================
 
!-----------------------------------------
! i1,j1 => integer. lower left grid point
! i2,j1 => integer. lower right grid point
! i1,j2 => integer. upper left grid point
! i2,j2 => integer. upper right grid point
!-----------------------------------------
   i1=bisec(xgrid,xm)
   i2=i1+1
   dxm=xm-xgrid(i1)   
   if(dxm.eq.0d0) i2=1
   
   j1=bisec(zgrid,zm)
   j2=j1+1
   dzm=zm-zgrid(j1)
   if(dzm.eq.0d0) j2=1
   
! Interpolate the field 'grdval' to the marker
   mkval=(1-dxm/delx(i2))*(1-dzm/delz(j2))*grdval(i1,j1) &
        +(dxm/delx(i2))*(1-dzm/delz(j2))*grdval(i2,j1)   &
        +(1-dxm/delx(i2))*(dzm/delz(j2))*grdval(i1,j2)   &
        +(dxm/delx(i2))*(dzm/delz(j2))*grdval(i2,j2)

 end subroutine grid2mk
 

!========================================================================
 function bisec(xx,x)
! Bisection algorithm to locate the position of 'x' on an 
! ascending grid contained in 'xx'
! On output the index 'jl' is returned such that xx(jl) <= x <= xx(jl+1)
! 
! Input:
! xx => real, array(0:). Either x- or z-grid
! x => real. x- or z- coord of the point to be located
! Output: 
! bisec => integer. Left index of the bounding cell
!========================================================================
   implicit none
   integer :: bisec,jl,ju,jm,n
   real(kind=8), dimension(0:), intent(in) :: xx
   real(kind=8), intent(in) :: x

   n=size(xx)-2
   jl=0
   ju=n+1
   do
     if(ju-jl.le.1) exit
     jm=(ju+jl)/2
     if(x.ge.xx(jm)) then
       jl=jm
     else
       ju=jm
     end if
   end do
   
   if(x.eq.xx(0)) then
     bisec=0
   elseif(x.eq.xx(n+1)) then
     bisec=n+1
   else
     bisec=jl
   end if
   
 end function bisec
 


 end program stream
